package com.anxuan.agent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.anxuan.agent.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 处理工具调用的基础代理类，具体实现 think 和 act 方法，可以用作创建实例的父类
 */

@EqualsAndHashCode(callSuper=true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent{

    //可用的工具
    private final ToolCallback[] availableTools;

    //保存工具调用信息的响应结果 （要调用哪些工具）
    private ChatResponse toolCallChatResponse;

    // 工具调用管理者
    private final ToolCallingManager toolCallingManager;

    //禁用Spring AI 内置的工具机制，自己维护选项和消息上下文
    private final ChatOptions chatOptions;


    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();

    }


    @Override
    public boolean think() {

        try {
            //1.校验提示词，拼接用户提示词
            if (StrUtil.isNotBlank(getNextStepPrompt())) {
                UserMessage userMessage = new UserMessage(getNextStepPrompt());
                getMessageList().add(userMessage);
            }
            //2.调用AI大模型，获取工具调用结果

            List<Message> messageList = getMessageList();
            Prompt prompt = new Prompt(messageList, this.chatOptions);
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();

            // 记录响应，用于等下Act
            this.toolCallChatResponse = chatResponse;
            //3.解析工具调用结果，获取要调用的结果
            //助手消息
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            //获取要调用的工具
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            //输出提示消息
            String result = assistantMessage.getText();
            log.info(getName(), "的思考: " + result);
            log.info(getName(), "选择了 " + toolCallList.size() + " 个工具来调用");

            String toolCallInfo = toolCallList.stream()
                    .map(toolCall -> {
                        return String.format("工具名称: %s, 参数: %s", toolCall.name(), toolCall.arguments());
                    })
                    .collect(Collectors.joining("\n"));
            log.info(toolCallInfo);
            //如果不需要调用工具，返回false
            if (toolCallList.isEmpty()) {
                //只有不调用工具是，才需要手动记录助手消息
                getMessageList().add(assistantMessage);
                return false;
            } else {
                // 需要调用是，无需记录助手消息，因为调用工具会自动记录
                return true;
            }
        }catch (Exception e){
            log.error(getName() + "的思考过程遇到了问题： ", e.getMessage());
            getMessageList().add(new AssistantMessage("处理时遇到了问题： " + e.getMessage()));
            return false;
        }
    }

    /**
     * 执行工具调用并处理结果
     * @return 执行结果
     */
    @Override
    public String act() {
        if(!toolCallChatResponse.hasToolCalls()){
            return "没有工具调用";
        }

        Prompt prompt = new Prompt(getMessageList(), this.chatOptions);
        //调用工具
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        //记录消息上下文,conversationHistory() 包含了工具调用和结果
        setMessageList(toolExecutionResult.conversationHistory());

        ToolResponseMessage toolResponseMessage = (ToolResponseMessage)CollUtil.getLast(toolExecutionResult.conversationHistory());

        //判断是否调用了终止工具
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> "doTerminate".equals(response.name()));
        if(terminateToolCalled){
            setState(AgentState.FINISHED);
        }
        String result = toolResponseMessage.getResponses().stream()
                .map(response -> "工具 " + response.name() + " 返回的结果: " + response.responseData())
                .collect(Collectors.joining("\n"));
        log.info(result);
        return result;
    }
}
